#!/usr/bin/env python3
"""
数据拆分插件测试脚本
测试DataSplitter插件的各种拆分策略
"""

import sys
import os

# 添加项目根目录到Python路径
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from src.plugins.data_splitter.data_splitter import DataSplitter


def create_test_data():
    """创建测试数据"""
    return [
        {"id": 1, "name": "Alice", "age": 25, "city": "北京", "salary": 5000, "department": "IT"},
        {"id": 2, "name": "Bob", "age": 30, "city": "上海", "salary": 6000, "department": "HR"},
        {"id": 3, "name": "Charlie", "age": 35, "city": "广州", "salary": 7000, "department": "Finance"},
        {"id": 4, "name": "David", "age": 28, "city": "深圳", "salary": 5500, "department": "IT"},
        {"id": 5, "name": "Eve", "age": 32, "city": "北京", "salary": 6500, "department": "HR"},
        {"id": 6, "name": "Frank", "age": 29, "city": "上海", "salary": 5800, "department": "Finance"},
        {"id": 7, "name": "Grace", "age": 26, "city": "广州", "salary": 5200, "department": "IT"},
        {"id": 8, "name": "Henry", "age": 31, "city": "深圳", "salary": 6200, "department": "HR"},
        {"id": 9, "name": "Ivy", "age": 27, "city": "北京", "salary": 5400, "department": "Finance"},
        {"id": 10, "name": "Jack", "age": 33, "city": "上海", "salary": 6800, "department": "IT"}
    ]


def test_ratio_split():
    """测试按比例拆分"""
    print("=== 测试按比例拆分 ===")
    
    # 创建插件实例
    config = {
        "config": {
            "split_strategy": "ratio",
            "ratios": [0.7, 0.3],  # 70%训练集，30%测试集
            "random_seed": 42,     # 固定随机种子保证可重现
            "shuffle": True        # 打乱数据
        }
    }
    plugin = DataSplitter("test_ratio_split", config)
    
    # 准备测试数据
    test_data = create_test_data()
    input_data = {"data_loader": {"raw_data": test_data}}
    
    # 执行拆分
    result = plugin.execute(input_data)
    
    # 验证结果
    if result["split_data"]:
        train_set, test_set = result["split_data"]
        print(f"训练集大小: {len(train_set)} 条记录")
        print(f"测试集大小: {len(test_set)} 条记录")
        print(f"拆分信息: {result['split_info']}")
        print("训练集前3条记录:")
        for i, record in enumerate(train_set[:3]):
            print(f"  {i+1}. {record}")
        print("测试集前3条记录:")
        for i, record in enumerate(test_set[:3]):
            print(f"  {i+1}. {record}")
    else:
        print(f"拆分失败: {result['error']}")
    
    print()


def test_condition_split():
    """测试按条件拆分"""
    print("=== 测试按条件拆分 ===")
    
    # 创建插件实例 - 按年龄拆分
    config = {
        "config": {
            "split_strategy": "condition",
            "condition_field": "age",
            "condition_value": 30,
            "condition_operator": ">="  # 年龄>=30的为一组，其他的为另一组
        }
    }
    plugin = DataSplitter("test_condition_split", config)
    
    # 准备测试数据
    test_data = create_test_data()
    input_data = {"data": test_data}  # 使用直接数据格式
    
    # 执行拆分
    result = plugin.execute(input_data)
    
    # 验证结果
    if result["split_data"]:
        senior_group, junior_group = result["split_data"]
        print(f"年龄>=30的组: {len(senior_group)} 条记录")
        print(f"年龄<30的组: {len(junior_group)} 条记录")
        print("年龄>=30的组:")
        for record in senior_group:
            print(f"  {record['name']}: {record['age']}岁")
        print("年龄<30的组:")
        for record in junior_group:
            print(f"  {record['name']}: {record['age']}岁")
    else:
        print(f"拆分失败: {result['error']}")
    
    print()


def test_field_value_split():
    """测试按字段值拆分"""
    print("=== 测试按字段值拆分 ===")
    
    # 创建插件实例 - 按部门拆分
    config = {
        "config": {
            "split_strategy": "field_value",
            "split_field": "department",
            "max_groups": 5  # 最多5个分组
        }
    }
    plugin = DataSplitter("test_field_value_split", config)
    
    # 准备测试数据
    test_data = create_test_data()
    input_data = test_data  # 使用列表格式
    
    # 执行拆分
    result = plugin.execute(input_data)
    
    # 验证结果
    if result["split_data"]:
        groups = result["split_data"]
        print(f"按部门拆分为 {len(groups)} 个组:")
        for i, group in enumerate(groups):
            if group:
                department = group[0]["department"]
                print(f"  {department}部门: {len(group)} 人")
                for record in group:
                    print(f"    - {record['name']}")
    else:
        print(f"拆分失败: {result['error']}")
    
    print()


def test_count_split():
    """测试按数量拆分"""
    print("=== 测试按数量拆分 ===")
    
    # 创建插件实例
    config = {
        "config": {
            "split_strategy": "count",
            "counts": [3, 3, 4],  # 拆分为3组，每组分别有3、3、4条记录
            "remainder_strategy": "distribute"  # 剩余数据分配到前几组
        }
    }
    plugin = DataSplitter("test_count_split", config)
    
    # 准备测试数据
    test_data = create_test_data()
    input_data = {"raw_data": test_data}  # 使用节点数据格式
    
    # 执行拆分
    result = plugin.execute(input_data)
    
    # 验证结果
    if result["split_data"]:
        groups = result["split_data"]
        print(f"拆分为 {len(groups)} 个组:")
        for i, group in enumerate(groups):
            print(f"  组{i+1}: {len(group)} 条记录")
            for j, record in enumerate(group[:2]):  # 显示前2条记录
                print(f"    {j+1}. {record['name']} (ID: {record['id']})")
            if len(group) > 2:
                print(f"    ... 还有 {len(group) - 2} 条记录")
    else:
        print(f"拆分失败: {result['error']}")
    
    print()


def test_invalid_config():
    """测试无效配置"""
    print("=== 测试无效配置 ===")
    
    # 测试1: 无效的拆分策略
    config1 = {"config": {"split_strategy": "invalid_strategy"}}
    plugin1 = DataSplitter("test_invalid_strategy", config1)
    result1 = plugin1.execute(create_test_data())
    print(f"无效策略测试结果: {result1['error']}")
    
    # 测试2: 缺少必需字段
    config2 = {"config": {"split_strategy": "condition"}}  # 缺少condition_field
    plugin2 = DataSplitter("test_missing_field", config2)
    result2 = plugin2.execute(create_test_data())
    print(f"缺少字段测试结果: {result2['error']}")
    
    # 测试3: 无效数据格式
    config3 = {"config": {"split_strategy": "ratio", "ratios": [0.5, 0.5]}}
    plugin3 = DataSplitter("test_invalid_data", config3)
    result3 = plugin3.execute("invalid_data")  # 字符串而不是列表
    print(f"无效数据格式测试结果: {result3['error']}")
    
    print()


def test_config_validation():
    """测试配置验证"""
    print("=== 测试配置验证 ===")
    
    # 有效配置
    valid_configs = [
        {"config": {"split_strategy": "ratio", "ratios": [0.7, 0.3]}},
        {"config": {"split_strategy": "condition", "condition_field": "age", "condition_value": 30}},
        {"config": {"split_strategy": "field_value", "split_field": "department"}},
        {"config": {"split_strategy": "count", "counts": [5, 5]}}
    ]
    
    print("有效配置验证:")
    for i, config in enumerate(valid_configs):
        plugin = DataSplitter(f"valid_config_{i}", config)
        result = plugin.validate_config()
        print(f"  配置{i+1}: {'通过' if result else '失败'}")
    
    # 无效配置
    invalid_configs = [
        {"config": {"split_strategy": "invalid"}},
        {"config": {"split_strategy": "ratio", "ratios": ["invalid"]}},
        {"config": {"split_strategy": "condition"}},  # 缺少condition_field
        {"config": {"split_strategy": "field_value"}},  # 缺少split_field
        {"config": {"split_strategy": "count", "counts": ["invalid"]}}
    ]
    
    print("\n无效配置验证:")
    for i, config in enumerate(invalid_configs):
        plugin = DataSplitter(f"invalid_config_{i}", config)
        result = plugin.validate_config()
        print(f"  配置{i+1}: {'通过' if result else '失败'}")
    
    print()


if __name__ == "__main__":
    print("开始数据拆分插件测试...\n")
    
    # 运行所有测试
    test_ratio_split()
    test_condition_split()
    test_field_value_split()
    test_count_split()
    test_invalid_config()
    test_config_validation()
    
    print("所有测试完成！")